package encrypt

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"os"
	"path/filepath"
	"sync"
)

var (
	onceRsa   sync.Once
	singleRsa *Rsa
)

type Rsa struct {
	// 证书大小
	bits        int
	privateFile string
	PrivateKey  *rsa.PrivateKey
	publicFile  string
	PublicKey   *rsa.PublicKey

	// 存储路径
	storageDir string
}

// GetRSA 获取RSA加密实例
func GetRSA() *Rsa {
	onceRsa.Do(func() {
		singleRsa = &Rsa{
			bits:        2048,
			privateFile: "private.pem",
			publicFile:  "public.pem",
		}
	})
	return singleRsa
}

func (r *Rsa) SetBits(size int) *Rsa {
	r.bits = size
	return r
}

// SetStorageDir 设置密钥保存目录
func (r *Rsa) SetStorageDir(dir string) *Rsa {
	r.storageDir = dir
	return r
}

func (r Rsa) getFilepath(filename string) string {
	return filepath.Join(r.storageDir, filename)
}

// GeneratePrivateKey 生成私钥
func (r *Rsa) GeneratePrivateKey() error {
	privateKey, err := rsa.GenerateKey(rand.Reader, r.bits)
	if err != nil {
		return err
	}
	r.PrivateKey = privateKey
	// 通过 x509标准将得到的ras私钥序列化为ASN.1的DER编码字符串
	x509PrivateKey := x509.MarshalPKCS1PrivateKey(r.PrivateKey)
	// 保存私钥
	file, err := os.Create(r.privateFile)
	if err != nil {
		return err
	}

	defer file.Close()
	// 构建一个pem.Block结构体对象
	block := pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509PrivateKey}
	// 保存到文件
	err = pem.Encode(file, &block)
	if err != nil {
		return err
	}
	return nil
}

// GeneratePublicKey 生成公钥
func (r *Rsa) GeneratePublicKey() error {
	publicKey := r.PrivateKey.PublicKey
	x509PublicKey, err := x509.MarshalPKIXPublicKey(&publicKey)
	if err != nil {
		return err
	}
	// 保存公钥到文件
	file, err := os.Create(r.getFilepath(r.publicFile))
	if err != nil {
		return err
	}
	defer file.Close()

	block := pem.Block{Type: "RSA PUBLIC KEY", Bytes: x509PublicKey}
	err = pem.Encode(file, &block)
	if err != nil {
		return err
	}

	return nil
}

func loadPemBlock(filepath string) (*pem.Block, error) {
	// 加载文件
	file, err := os.Open(filepath)
	if err != nil {
		return nil, err
	}

	defer file.Close()
	info, _ := file.Stat()
	buf := make([]byte, info.Size())
	// pem解码
	block, _ := pem.Decode(buf)
	return block, nil
}

// LoadPrivateKey 加密私钥
func (r *Rsa) LoadPrivateKey(filename string) error {
	if r.PrivateKey != nil {
		return nil
	}
	block, err := loadPemBlock(r.getFilepath(filename))
	if err != nil {
		return err
	}
	// x509解码
	r.PrivateKey, err = x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		return err
	}

	return nil
}

func (r *Rsa) LoadPublicKey(filename string) error {
	if r.PublicKey != nil {
		return nil
	}
	// 加载文件
	block, err := loadPemBlock(r.getFilepath(filename))
	if err != nil {
		return err
	}

	r.PublicKey, err = x509.ParsePKCS1PublicKey(block.Bytes)
	if err != nil {
		return err
	}
	return nil
}

// EncryptPublic 公钥加密
func (r Rsa) EncryptPublic(cipherText []byte) ([]byte, error) {
	return rsa.EncryptPKCS1v15(rand.Reader, r.PublicKey, cipherText)
}

// DecryptPrivate 私钥解密
func (r Rsa) DecryptPrivate(cipherText []byte) ([]byte, error) {
	return rsa.DecryptPKCS1v15(rand.Reader, r.PrivateKey, cipherText)
}
